// SPDX-FileCopyrightText: 2023-2025 Sangfor Technologies Inc.
// SPDX-License-Identifier: Mulan PSL v2
package com.geniusai.aip.auth.service.impl;

import com.geniusai.aip.auth.config.ExternalConfig;
import com.geniusai.aip.auth.constants.GeniusVGpuEnum;
import com.geniusai.aip.auth.constants.ResourceConstants;
import com.geniusai.aip.auth.domain.GpuResponse;
import com.geniusai.aip.auth.domain.Notebooks;
import com.geniusai.aip.auth.domain.TotalResource;
import com.geniusai.aip.auth.domain.UserDO;
import com.geniusai.aip.auth.domain.VGpuAllocateVO;
import com.geniusai.aip.auth.mapper.NotebooksMapper;
import com.geniusai.aip.auth.mapper.PvcsMapper;
import com.geniusai.aip.auth.mapper.UserMapper;
import com.geniusai.aip.auth.service.MetricService;
import com.geniusai.aip.feature.common.web.utils.UrlUtils;
import io.fabric8.kubernetes.api.model.Node;
import io.fabric8.kubernetes.api.model.NodeCondition;
import io.fabric8.kubernetes.api.model.PersistentVolumeClaim;
import io.fabric8.kubernetes.api.model.Quantity;
import io.fabric8.kubernetes.api.model.ResourceQuota;
import io.fabric8.kubernetes.api.model.metrics.v1beta1.ContainerMetrics;
import io.fabric8.kubernetes.api.model.metrics.v1beta1.NodeMetrics;
import io.fabric8.kubernetes.api.model.metrics.v1beta1.PodMetrics;
import io.fabric8.kubernetes.client.KubernetesClient;
import io.fabric8.kubernetes.client.dsl.ExecListener;
import io.fabric8.kubernetes.client.dsl.ExecWatch;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.web.client.RestTemplate;

import java.io.ByteArrayOutputStream;
import java.io.OutputStream;
import java.math.BigDecimal;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

/**
 * @author hml14563
 * @date 2021/5/6
 */
@Service
@Slf4j
public class MetricServiceImpl implements MetricService {
    @Autowired
    private KubernetesClient client;

    @Autowired
    private ExternalConfig externalConfig;

    @Autowired
    private UserMapper userMapper;

    @Autowired
    private NotebooksMapper notebooksMapper;

    @Autowired
    private PvcsMapper pvcsMapper;

    @Autowired
    private RestTemplate restTemplate;

    private static final String QUERY_GPU_URL = "api/v1/query?query={query}&time={time}";

    private static final String MONITORING = "monitoring";

    private static final String ZERO = "0";

    @Override
    public Map<String, Double> getUsedResourceByNamespace(String namespace) {
        HashMap<String, Double> resourceMap = new HashMap<>(2);
        double cpuNum = 0, memorySize = 0;
        for (PodMetrics pod : client.top().pods().metrics(namespace).getItems()) {
            for (ContainerMetrics container : pod.getContainers()) {
                // 获取已使用cpu大小
                cpuNum = getUsedResource(container, cpuNum, ResourceConstants.CPU);
                resourceMap.put(ResourceConstants.CPU, cpuNum);
                // 获取已使用内存大小
                memorySize = getUsedResource(container, memorySize, ResourceConstants.MEMORY);
                resourceMap.put(ResourceConstants.MEMORY, memorySize);
            }
        }
        return resourceMap;
    }

    @Override
    public Map<String, Double> getRequestResourceByNamespace(String namespace) {
        HashMap<String, Double> requestResourceMap = new HashMap<>(2);
        double cpuNum = 0, memorySize = 0;
        for (ResourceQuota resourceQuota : client.resourceQuotas().inNamespace(namespace).list().getItems()) {
            // 获取分配cpu大小
            cpuNum = getRequestResource(resourceQuota, cpuNum, ResourceConstants.LIMITS_CPU);
            requestResourceMap.put(ResourceConstants.LIMITS_CPU, cpuNum);
            // 获取分配内存大小
            memorySize = getRequestResource(resourceQuota, memorySize, ResourceConstants.LIMITS_MEMORY);
            requestResourceMap.put(ResourceConstants.LIMITS_MEMORY, memorySize);
        }
        return requestResourceMap;
    }

    @Override
    public Map<String, Double> getTotalResource() {
        HashMap<String, Double> resourceMap = new HashMap<>(8);
        double totalCpu = 0, totalUsedCpu = 0, totalGpu = 0, totalMemory = 0, totalUsedMemory = 0;
        for (Node node : client.nodes().list().getItems()) {
            log.info("节点名称：{}", node.getMetadata().getName());
            List<NodeCondition> conditions = node.getStatus().getConditions();
            log.info("节点状态:{}", conditions.get(conditions.size() - 1).getStatus());
            // 先判断节点的状态信息
            if ("True".equals(conditions.get(conditions.size() - 1).getStatus())) {
                if ("true".equals(node.getMetadata().getLabels().get(ResourceConstants.WORKER_ROLE))) {
                    // 获取集群的cpu资源
                    totalCpu = getTotalResource(node, totalCpu, ResourceConstants.CPU);
                    resourceMap.put(ResourceConstants.CPU, totalCpu);
                    // 获取集群的memory资源
                    totalMemory = getTotalResource(node, totalMemory, ResourceConstants.MEMORY);
                    log.info("总内存totalMemory为：{}", totalMemory);
                    resourceMap.put(ResourceConstants.MEMORY, totalMemory);

                    NodeMetrics nodeMetrics = client.top().nodes().metrics(node.getMetadata().getName());
                    // 获取集群已使用的cpu资源
                    totalUsedCpu = getUsedTotalResource(nodeMetrics, totalUsedCpu, ResourceConstants.CPU);
                    resourceMap.put(ResourceConstants.USED_CPU, totalUsedCpu);

                    // 获取集群已使用的memory资源
                    totalUsedMemory = getUsedTotalResource(nodeMetrics, totalUsedMemory, ResourceConstants.MEMORY);
                    resourceMap.put(ResourceConstants.USED_MEMORY, totalUsedMemory);
                    log.info("已使用内存totalUsedMemory为：{}", totalMemory);
                }
                if (node.getStatus().getAllocatable().containsKey(ResourceConstants.NVIDIAGPU) && "true".equals(node.getMetadata().getLabels().get(ResourceConstants.WORKER_ROLE))) {
                    // 获取集群的gpu资源
                    totalGpu = totalGpu + Integer.parseInt(node.getStatus().getAllocatable().get(ResourceConstants.NVIDIAGPU).getAmount());
                    resourceMap.put(ResourceConstants.GPU, totalGpu);
                }
            }
        }
        return resourceMap;
    }

    private Double getTotalResource(Node node, Double num, String type) {
        Quantity quantity = node.getStatus().getAllocatable().get(type);
        return getResourceSize(num, type, quantity);
    }

    private Double getUsedTotalResource(NodeMetrics nodeMetrics, Double num, String type) {
        Quantity quantity = nodeMetrics.getUsage().get(type);
        return getResourceSize(num, type, quantity);
    }

    private Double getRequestResource(ResourceQuota resourceQuota, Double num, String type) {
        Quantity quantity = resourceQuota.getStatus().getUsed().get(type);
        return getResourceSize(num, type, quantity);
    }

    private Double getUsedResource(ContainerMetrics containerMetrics, Double num, String type) {
        Quantity quantity = containerMetrics.getUsage().get(type);
        return getResourceSize(num, type, quantity);
    }

    private Double getResourceSize(Double num, String type, Quantity quantity) {
        if (quantity != null) {
            if (type.equals(ResourceConstants.CPU) || type.equals(ResourceConstants.LIMITS_CPU)) {
                num = num + CpuNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
            } else {
                num = num + MemoryNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
            }
        }
        return num;
    }

    @Override
    public Map<String, Double> getRequestTotalResource() {
        HashMap<String, Double> requestResourceMap = new HashMap<>(4);
        double totalRequestCpu = 0, totalRequestGpu = 0, totalRequestMemory = 0;
        List<UserDO> users = userMapper.selectList(null);
        if (!Objects.isNull(users)) {
            for (UserDO user : users) {
                totalRequestCpu = totalRequestCpu + user.getCpuNum();
                totalRequestMemory = totalRequestMemory + user.getMemorySize();
            }
            requestResourceMap.put(ResourceConstants.TOTAL_REQUEST_CPU, totalRequestCpu);
            requestResourceMap.put(ResourceConstants.TOTAL_REQUEST_GPU, totalRequestGpu);
            requestResourceMap.put(ResourceConstants.TOTAL_REQUEST_MEMORY, totalRequestMemory);
        }
        return requestResourceMap;
    }

    @Override
    public Double getCpuByNamespace(String namespace) {
        double cpuNum = 0;
        for (PodMetrics pod : client.top().pods().metrics(namespace).getItems()) {
            for (ContainerMetrics container : pod.getContainers()) {
                Quantity quantity = container.getUsage().get(ResourceConstants.CPU);
                if (quantity != null) {
                    cpuNum = cpuNum + CpuNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
                }
            }
        }
        return cpuNum;
    }

    @Override
    public Double getRequestCpuByNamespace(String namespace) {
        double cpuNum = 0;
        for (ResourceQuota resourceQuota : client.resourceQuotas().inNamespace(namespace).list().getItems()) {
            Quantity quantity = resourceQuota.getStatus().getUsed().get(ResourceConstants.LIMITS_CPU);
            if (quantity != null) {
                cpuNum = cpuNum + CpuNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
            }
        }
        return cpuNum;
    }

    @Override
    public Integer getGpuByNamespace(String namespace) {
        int gpuNum = 0;
        GpuResponse response = getGpuResponse();
        if (response.getData().getResult().size() == 0) {
            return gpuNum;
        }
        for (GpuResponse.GpuMetric gpuMetric : response.getData().getResult()) {
            if (!ZERO.equals(gpuMetric.getValue().get(1)) && namespace.equals(gpuMetric.getMetric().getExportedNamespace())) {
                gpuNum++;
            }
        }
        return gpuNum;
    }

    @Override
    public Integer getRequestGpuByNamespace(String namespace) {
        int gpuNum = 0;
        GpuResponse response = getGpuResponse();
        if (response.getData().getResult() == null) {
            return gpuNum;
        }
        for (GpuResponse.GpuMetric gpuMetric : response.getData().getResult()) {
            if (gpuMetric.getMetric().getExportedNamespace() != null && gpuMetric.getMetric().getExportedNamespace().equals(namespace)) {
                gpuNum++;
            }
        }
        return gpuNum;
    }

    @Override
    public Double getMemoryByNamespace(String namespace) {
        double memorySize = 0;
        for (PodMetrics pod : client.top().pods().metrics(namespace).getItems()) {
            for (ContainerMetrics container : pod.getContainers()) {
                Quantity quantity = container.getUsage().get(ResourceConstants.MEMORY);
                if (quantity != null) {
                    memorySize = memorySize + MemoryNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
                }
            }
        }
        return memorySize;
    }

    @Override
    public Double getRequestMemoryByNamespace(String namespace) {
        double memorySize = 0;
        for (ResourceQuota resourceQuota : client.resourceQuotas().inNamespace(namespace).list().getItems()) {
            Quantity quantity = resourceQuota.getStatus().getUsed().get(ResourceConstants.LIMITS_MEMORY);
            if (quantity != null) {
                memorySize = memorySize + MemoryNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
            }
        }
        return memorySize;
    }

    @Override
    public Double getStorageByNamespace(String namespace) {
        double storageSize = 0;
        for (PersistentVolumeClaim pvc : client.persistentVolumeClaims().inNamespace(namespace).list().getItems()) {
            Map<String, Quantity> map = pvc.getSpec().getResources().getRequests();
            if (map != null) {
                storageSize = storageSize + MemoryNumEnum.getFormat(map.get(ResourceConstants.STORAGE).getFormat())
                        * Double.parseDouble(map.get(ResourceConstants.STORAGE).getAmount());
            }
        }
        return storageSize;
    }

    @Override
    public Double getTotalCpu() {
        double totalCpu = 0;
        for (Node node : client.nodes().list().getItems()) {
            if ("true".equals(node.getMetadata().getLabels().get(ResourceConstants.WORKER_ROLE))) {
                Quantity quantity = node.getStatus().getAllocatable().get(ResourceConstants.CPU);
                totalCpu = totalCpu + CpuNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
            }
        }
        return totalCpu;
    }

    @Override
    public Double getTotalUsedCpu() {
        double totalUsedCpu = 0;
        for (Node node : client.nodes().list().getItems()) {
            if ("true".equals(node.getMetadata().getLabels().get(ResourceConstants.WORKER_ROLE))) {
                NodeMetrics nodeMetrics = client.top().nodes().metrics(node.getMetadata().getName());
                Quantity quantity = nodeMetrics.getUsage().get(ResourceConstants.CPU);
                totalUsedCpu = totalUsedCpu + CpuNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
            }
        }
        return totalUsedCpu;
    }

    /**
     * add all the cpu num in database
     */
    @Override
    public Double getRequestTotalCpu() {
        double totalRequestCpu = 0;
        List<UserDO> users = userMapper.selectList(null);
        if (users == null) {
            return totalRequestCpu;
        }
        for (UserDO user : users) {
            totalRequestCpu = totalRequestCpu + user.getCpuNum();
        }
        return totalRequestCpu;
    }

    @Override
    public Integer getTotalGpu() {
        int totalGpu = 0;
        for (Node node : client.nodes().list().getItems()) {
            if (node.getStatus().getAllocatable().containsKey(ResourceConstants.NVIDIAGPU) && "true".equals(node.getMetadata().getLabels().get(ResourceConstants.WORKER_ROLE))) {
                totalGpu = totalGpu + Integer.parseInt(node.getStatus().getAllocatable().get(ResourceConstants.NVIDIAGPU).getAmount());
            }
        }
        return totalGpu;
    }

    @Override
    public Integer getTotalUsedGpu() {
        int totalUsedGpu = 0;
        GpuResponse response = getGpuResponse();
        if (response.getData().getResult() == null) {
            return totalUsedGpu;
        }
        for (GpuResponse.GpuMetric gpuMetric : response.getData().getResult()) {
            if (!ZERO.equals(gpuMetric.getValue().get(1))) {
                totalUsedGpu++;
            }
        }
        return totalUsedGpu;
    }

    @Override
    public Double getTotalMemory() {
        double totalMemory = 0;
        for (Node node : client.nodes().list().getItems()) {
            if ("true".equals(node.getMetadata().getLabels().get(ResourceConstants.WORKER_ROLE))) {
                Quantity quantity = node.getStatus().getAllocatable().get(ResourceConstants.MEMORY);
                totalMemory = totalMemory
                        + MemoryNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
            }
        }
        return totalMemory;
    }

    @Override
    public Double getTotalUsedMemory() {
        double totalUsedMemory = 0;
        for (Node node : client.nodes().list().getItems()) {
            if ("true".equals(node.getMetadata().getLabels().get(ResourceConstants.WORKER_ROLE))) {
                NodeMetrics nodeMetrics = client.top().nodes().metrics(node.getMetadata().getName());
                Quantity quantity = nodeMetrics.getUsage().get(ResourceConstants.MEMORY);
                totalUsedMemory = totalUsedMemory + MemoryNumEnum.getFormat(quantity.getFormat()) * Double.parseDouble(quantity.getAmount());
            }
        }
        return totalUsedMemory;
    }

    @Override
    public Double getRequestTotalMemory() {
        double totalRequestMemory = 0;
        List<UserDO> users = userMapper.selectList(null);
        if (users == null) {
            return totalRequestMemory;
        }
        for (UserDO user : users) {
            totalRequestMemory = totalRequestMemory + user.getMemorySize();
        }
        return totalRequestMemory;
    }

    @Override
    public VGpuAllocateVO getUserAllocateVGpu(String userId) {
        VGpuAllocateVO vGpuAllocateVO = new VGpuAllocateVO();
        List<Notebooks> notebooks = notebooksMapper.selectNotebooks(userId);
        if (!CollectionUtils.isEmpty(notebooks)) {
            int allocatedVGpuNum = notebooks.stream().filter(obj -> obj.getVGpuNum() != null).mapToInt(Notebooks::getVGpuNum).sum();
            int allocatedVGpuMemory = notebooks.stream().filter(obj -> obj.getVGpuMemory() != null).mapToInt(Notebooks::getVGpuMemory).sum();
            double allocatedMemorySizeDouble = notebooks.stream().filter(obj -> obj.getMemorySize() != null).mapToDouble(Notebooks::getMemorySize).sum();
            int allocatedMemorySize = new BigDecimal(allocatedMemorySizeDouble).setScale(0, BigDecimal.ROUND_HALF_DOWN).intValue();
            int allocatedCpuNum = notebooks.stream().filter(obj -> obj.getCpuNum() != null).mapToInt(Notebooks::getCpuNum).sum();

            vGpuAllocateVO.setAllocatedVGpuNum(allocatedVGpuNum);
            vGpuAllocateVO.setAllocatedVGpuMemory(allocatedVGpuMemory);
            vGpuAllocateVO.setAllocatedMemorySize(allocatedMemorySize);
            vGpuAllocateVO.setAllocatedCpuNum(allocatedCpuNum);
        }

        return vGpuAllocateVO;
    }

    @Override
    public List<TotalResource.VGpuGraph> getTotalVGpuMemoryAllocateUsage() {
        List<TotalResource.VGpuGraph> vGpuGraphs = new ArrayList<>();
        Map<String, TotalResource.VGpuGraph> vGpuAllocatedGraphs = notebooksMapper.selectVGpuMemoryGroupByVGpuDisplayName();
        for (GeniusVGpuEnum geniusVGpuEnum: GeniusVGpuEnum.values()) {
            TotalResource.VGpuGraph vGpuGraph = new TotalResource.VGpuGraph();
            vGpuGraph.setName((GeniusVGpuEnum.getSize(geniusVGpuEnum.getValue()) + "G"));
            vGpuGraph.setValue(vGpuAllocatedGraphs.getOrDefault(geniusVGpuEnum.getValue(), new TotalResource.VGpuGraph(0)).getValue());
            vGpuGraphs.add(vGpuGraph);
        }
        return vGpuGraphs;
    }

    @Override
    public Integer getAllocatedTotalVGpuMemory() {
        Integer vGpuMemorySum = userMapper.sumVGpuMemory();
        return vGpuMemorySum == null ? 0 : vGpuMemorySum;
    }

    @Override
    public Integer getAllocatedStorageByUserId(String userId) {
        Integer userPvcSize =  pvcsMapper.sumUserPvcSize(userId);
        return userPvcSize == null ? 0 : userPvcSize;
    }

    @Override
    public Double getTotalStorage() {
        return 0d;
    }

    @Override
    public Double getTotalUsedStorage() {
        return 0d;
    }

    private GpuResponse getGpuResponse() {
        Map<String, String> params = new HashMap<>(4);
        params.put("query", "DCGM_FI_DEV_GPU_UTIL");
        params.put("time", String.valueOf(new Timestamp(System.currentTimeMillis()).getTime() / 1000));
        return restTemplate.getForObject(
                UrlUtils.createUrl(
                        externalConfig.getPrometheusService(),
                        externalConfig.getPrometheusPort(), QUERY_GPU_URL
                ), GpuResponse.class, params
        );
    }
}
